#!/usr/bin/env python

""" Transform a coordinated mesh, represented by a .node file
containing the coordinates of the nodes and an .ele file containing
the description of its elements, into a nodal (files .nnode and
.ngraph) and dual (.dnode and .dgraph) graphs. Similar to mesh2nodal 
and mesh2dual Metis tools, they also output the coordinates of the 
resulting graphs.


WARNING: to respect
the metis convention (and results and partitioners like Pagrid), 
both elements and nodes numbering start from __1__!!!


Usage: ./mesh2graph.py [-n nodal graph] [-d dual graph] .node_file .ele_file

Format:
.node:
First line: #ofnodes dimension
ith line = coordinates of ie node: x y [z]

.mesh:
First line: #ofelements #ofnodes_in_one_element (3=triangle)
ith line = description of the ie element: node1 node2 node3

.nnode: (nodal node)
1st line: #ofnodes #dimension
ith line: coordinates of nnode i

.dnode: (dual node)
1st line: #ofnodes #dimension
ith line: coordinates of dnode i

.ngraph: (nodal graph connectivity)
1st line: #ofnodes #ofedges
ie line=nodes directly connected to node i: node1 node2 node3 node4 ...

.dgraph: (dual graph connectivity)
1st line: #ofelements #ofedges
ie line(for triangular elements): node1 node2 node3

For the Wrekagrid project,
Basile Clout, November 2007
"""



import os
import sys
import pprint
import getopt


def bary3(tuple1, tuple2, tuple3):
    """ Return the coordinates of the barycentre of the 3 given points
    """

    x = (tuple1[0]+tuple2[0]+tuple3[0])/3
    y = (tuple1[1]+tuple2[1]+tuple3[1])/3
    z = (tuple1[2]+tuple2[2]+tuple3[2])/3
    
    return (x, y, z)

def bary2(tuple1, tuple2, tuple3):
    """ Return the coordinates of the barycentre of the 3 given points
    """

    x = (tuple1[0]+tuple2[0]+tuple3[0])/3
    y = (tuple1[1]+tuple2[1]+tuple3[1])/3
    
    return (x, y)
    

def main():

    in_node = "test.node"
    in_mesh = "test.mesh"
    out_ngraph = "test.ngraph"
    out_nnode = "test.node.ngraph"
    out_dgraph = "test.dgraph"
    out_dnode = "test.node.dgraph"

    fl=0
    debug=0
    
    try:
        opts, arg = getopt.getopt(sys.argv[1:], "ndh", ["debug"])
    except getopt.error, msg:
        print msg
        print "For help, use the -h switch"
        sys.exit(2)

    for o, a in opts:

        if o=="-h":
            print __doc__
            sys.exit(0)
        elif o=="-d":
            fl=2
        elif o=="-n":
            fl=1
        elif o=="--debug":
            debug=1
        else:
            print "Unknown option. bye"
            sys.exit(2)

    if len(arg)!=2:
        print "Please enter a .node and .mesh/.ele file"
        sys.exit(2)
    else:
        in_node=arg[0]
        in_mesh=arg[1]
        prefix = in_mesh.rsplit('.',1)[0]
        out_ngraph = prefix+".ngraph"
        out_nnode = prefix+".nnode"
        out_dgraph = prefix+".dgraph"
        out_dnode = prefix+".dnode"

    try:
        fnode = file(in_node, "r")
    except IOError:
        print "Impossible to open "+in_node
        sys.exit(2)

    try:
        fmesh = file(in_mesh, "r")
    except IOError:
        print "Impossible to open "+in_mesh
        sys.exit(2)


    nodes=[]
    elements=[]
    nbnodes=0
    nbelements=0
    dim=3
    nbpts=0
    cnt=-1

  # Read input files
    for line in fnode:

        if line[0] == "#":
            continue
        cnt+=1

        parts = line.split()            
        if cnt==0:
            (nbnodes, dim)=(int(parts[0]), int(parts[1]))
        else:
            nodes.append(tuple([float(x) for x in parts[1:4]]))

    if len(nodes)!=nbnodes:
        print "ERROR parsing node file ", str(nbnodes), str(len(nodes))
        sys.exit(2)
        
    cnt=-1
    for line in fmesh:
        if line[0] == "#":
            continue
        cnt+=1
        parts = line.split()
        if cnt==0:
            (nbelements, nbpts)=(int(parts[0]), int(parts[1]))
        else:
            elements.append(tuple([int(x) for x in parts]))

    if len(elements)!=nbelements:
        print "ERROR parsing node file ", str(nbelements), str(len(elements))
        sys.exit(2)
        
    if debug==1:
        print "Input nodes: ", pprint.pprint(nodes)
        print "Input elements: ", pprint.pprint(elements)


    if fl in (0,1):
        try:
            fnnode = file(out_nnode, "w")
        except IOError:
            print "Impossible to create "+in_node
            sys.exit(2)

        try:
            fngraph = file(out_ngraph, "w")
        except IOError:
            print "Impossible to create "+in_mesh
            sys.exit(2)

        nnode=[]
        ngraph=[]

        # Write Nodal graph
        nbnodes = len(nodes)
        nbelements = len(elements)
        links=0

        for n in range(0, nbnodes):
            ngraph.append([])
            nnode.append(nodes[n])
            for e in range(0,nbelements):
                if n+1 in elements[e]:
                    for a in elements[e]:
                        if a not in ngraph[n] and a != n+1:
                            ngraph[n].append(a)
                            links+=1

        if debug==1:
            print "Ngraph: ", pprint.pprint(ngraph)
            print "Nnode: ", pprint.pprint(nnode)
            print "Number of links ", links
            
        fnnode.write("# Nnode file generated from %s and %s\n" % (in_node, in_mesh))
        fnnode.write("%d %d\n" % (nbnodes, dim))
        for n in range(0, nbnodes):
            fnnode.write(str(n+1)+" %f %f %f\n" % nnode[n])
        print out_nnode+" successfully written!"
        
#        fngraph.write("# Ngraph file generated from %s and %s\n" % (in_node, in_mesh))
        fngraph.write("%d %d\n" % (nbnodes, links/2))
        for n in range(0, nbnodes):
            for i in ngraph[n]:
                fngraph.write(" %d" % i)
            fngraph.write("\n")
        print out_ngraph+" successfully written!"


    if fl in (0,2):
        try:
            fdnode = file(out_dnode, "w")
        except IOError:
            print "Impossible to create "+in_mesh
            sys.exit(2)

        try:
            fdgraph = file(out_dgraph, "w")
        except IOError:
            print "Impossible to create "+in_node
            sys.exit(2)


        dgraph=[]
        dnode=[]
        links=0

        # Write dual graph
        for e in range(0, nbelements):
            pt1 = elements[e][0]
            pt2 = elements[e][1]
            pt3 = elements[e][2]
            
            dnode.append(bary3(nodes[pt1-1], nodes[pt2-1], nodes[pt3-1]))
            dgraph.append([])
            for e2 in range(0, nbelements):
                if e!=e2:
                    ei = elements[e2]
                    if ((pt1 in ei and pt2 in ei) or (pt1 in ei and pt3 in ei) or (pt2 in ei and pt3 in ei)):
                        dgraph[e].append(e2+1)
                        links+=1

        if debug==1:
            print "Dgraph: ", pprint.pprint(dgraph)
            print "dnode: ", pprint.pprint(dnode)
        
        fdnode.write("# Dnode file generated from %s and %s\n" % (in_node, in_mesh))
        fdnode.write("%d %d\n" % (nbelements, dim))
        for n in range(0, nbelements):
            fdnode.write(str(n+1)+" %f %f %f\n" % dnode[n])
        print out_dnode+" successfully written!"
        
 #       fdgraph.write("# Dgraph file generated from %s and %s\n" % (in_node, in_mesh))
        fdgraph.write("%d %d\n" % (nbelements, links/2))
        for n in range(0, nbelements):
            for i in dgraph[n]:
                fdgraph.write(" %d" % i)
            fdgraph.write("\n")
        print out_dgraph+" successfully written!"

if __name__=="__main__":
    main()
